// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/common/base/status.h>
#include <pollux/serializers/presto_header.h>

namespace kumo::pollux::serializer::presto {
    enum class TokenType {
        HEADER,
        NUM_COLUMNS,
        COLUMN_ENCODING,
        NUM_ROWS,
        NULLS,
        BYTE_ARRAY,
        SHORT_ARRAY,
        INT_ARRAY,
        LONG_ARRAY,
        INT128_ARRAY,
        VARIABLE_WIDTH_DATA_SIZE,
        VARIABLE_WIDTH_DATA,
        DICTIONARY_INDICES,
        DICTIONARY_ID,
        HASH_TABLE_SIZE,
        HASH_TABLE,
        NUM_FIELDS,
        OFFSETS,
    };

    struct Token {
        TokenType tokenType;
        uint32_t length;
    };

    namespace detail {
        class PrestoVectorLexer {
        public:
            explicit PrestoVectorLexer(std::string_view source)
                : source_(source), committedPtr_(source.begin()) {
            }

            Status lex(std::vector<Token> &out) &&;

        private:
            Status lexHeader();

            Status lexColumEncoding(std::string &out);

            Status lexColumn();

            template<typename T>
            Status lexFixedArray(TokenType tokenType) {
                int32_t numRows;
                POLLUX_RETURN_NOT_OK(lexInt(TokenType::NUM_ROWS, &numRows));
                POLLUX_RETURN_NOT_OK(lexNulls(numRows));
                const auto numBytes = numRows * sizeof(T);
                POLLUX_RETURN_NOT_OK(lexBytes(numBytes, tokenType));
                return Status::OK();
            }

            Status lexVariableWidth();

            Status lexArray();

            Status lexMap();

            Status lexRow();

            Status lexDictionary();

            Status lexRLE();

            Status lexNulls(int32_t &numRows);

            Status lexBytes(int32_t numBytes, TokenType tokenType, char *dst = nullptr);

            template<typename T>
            Status lexInt(TokenType tokenType, T *out = nullptr) {
                assertCommitted();
                POLLUX_RETURN_IF(
                    source_.size() < sizeof(T),
                    Status::Invalid(
                        "Source size less than int size: {} < {}",
                        source_.size(),
                        sizeof(T)));
                const auto value = PrestoHeader::readInt<T>(&source_);
                if (out != nullptr) {
                    *out = value;
                }
                commit(tokenType);
                return Status::OK();
            }

            void assertCommitted() const {
                assert(committedPtr_ == source_.begin());
            }

            void commit(TokenType tokenType);

            std::string_view source_;
            const char *committedPtr_;
            std::vector<uint64_t> nullsBuffer_;
            std::vector<Token> tokens_;
        };
    } // namespace detail
} // namespace kumo::pollux::serializer::presto
